package com.example.ai.store;

import com.example.ai.bean.ChatMessages;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageDeserializer;
import dev.langchain4j.data.message.ChatMessageSerializer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;
import org.springframework.stereotype.Component;

import java.util.LinkedList;
import java.util.List;

/**
 * MongoDB聊天记录存储实现
 * 负责聊天消息的持久化管理，支持增删查操作
 *
 * @author czwa
 */
@Component
public class MongoChatMemoryStore implements ChatMemoryStore {
    @Autowired
    private MongoTemplate mongoTemplate;

    /**
     * 根据会话ID获取聊天消息列表
     *
     * @param memoryId 会话唯一标识
     * @return 聊天消息列表，若不存在则返回空列表
     */
    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        // 构建查询条件：根据memoryId查找文档
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);

        // 执行查询获取聊天消息文档
        ChatMessages chatMessages = mongoTemplate.findOne(query, ChatMessages.class);

        // 若文档不存在，返回空列表；否则解析JSON内容为消息列表
        return chatMessages == null
                ? new LinkedList<>()
                : ChatMessageDeserializer.messagesFromJson(chatMessages.getContent());
    }

    /**
     * 更新会话的聊天消息列表
     * 若会话不存在则创建新记录，存在则覆盖原有内容
     *
     * @param memoryId 会话唯一标识
     * @param messages 待保存的聊天消息列表
     */
    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        // 构建查询条件：根据memoryId查找文档
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);

        // 将消息列表序列化为JSON字符串并设置为更新内容
        Update update = new Update();
        update.set("content", ChatMessageSerializer.messagesToJson(messages));

        // 执行upsert操作：存在则更新，不存在则插入
        mongoTemplate.upsert(query, update, ChatMessages.class);
    }

    /**
     * 根据会话ID删除聊天记录
     *
     * @param memoryId 会话唯一标识
     */
    @Override
    public void deleteMessages(Object memoryId) {
        // 构建查询条件：根据memoryId查找文档
        Criteria criteria = Criteria.where("memoryId").is(memoryId);
        Query query = new Query(criteria);

        // 执行删除操作
        mongoTemplate.remove(query, ChatMessages.class);
    }
}